Two city scheduling [Quick Select]

Time: O(N) on average; Space: O(1); easy

There are 2N people a company is planning to interview.

The cost of flying the i-th person to city A is costs[i][0], and the cost of flying the i-th person to city B is costs[i][1].

Return the minimum cost to fly every person to a city such that exactly N people arrive in each city.

Example 1:

Input: costs = [[10,20], [30,200], [400,50], [30,20]]

Output: 110

Explanation:

  • The first person goes to city A for a cost of 10.

  • The second person goes to city A for a cost of 30.

  • The third person goes to city B for a cost of 50.

  • The fourth person goes to city B for a cost of 20.

The total minimum cost is 10 + 30 + 50 + 20 = 110 to have half the people interviewing in each city.

Notes:

  • 1 <= len(costs) <= 100

  • It is guaranteed that len(costs) is even.

  • 1 <= costs[i][0], costs[i][1] <= 1000

[1]:

import random
# quick select solution
class Solution1(object):
    """
    Time: O(N)~O(N^2), O(N) on average.
    Space: O(1)
    """
    def twoCitySchedCost(self, costs):
        """
        :type costs: List[List[int]]
        :rtype: int
        """
        def kthElement(nums, k, compare):
            def PartitionAroundPivot(left, right, pivot_idx, nums, compare):
                new_pivot_idx = left
                nums[pivot_idx], nums[right] = nums[right], nums[pivot_idx]
                for i in range(left, right):
                    if compare(nums[i], nums[right]):
                        nums[i], nums[new_pivot_idx] = nums[new_pivot_idx], nums[i]
                        new_pivot_idx += 1

                nums[right], nums[new_pivot_idx] = nums[new_pivot_idx], nums[right]
                return new_pivot_idx

            left, right = 0, len(nums) - 1
            while left <= right:
                pivot_idx = random.randint(left, right)
                new_pivot_idx = PartitionAroundPivot(left, right, pivot_idx, nums, compare)
                if new_pivot_idx == k:
                    return
                elif new_pivot_idx > k:
                    right = new_pivot_idx - 1
                else:  # new_pivot_idx < k.
                    left = new_pivot_idx + 1

        kthElement(costs, len(costs)//2, lambda a, b: a[0]-a[1] < b[0]-b[1])
        result = 0
        for i in range(len(costs)):
            result += costs[i][0] if i < len(costs)//2 else costs[i][1]
        return result

[3]:

s = Solution1()
costs =  [[10,20], [30,200], [400,50], [30,20]]
assert s.twoCitySchedCost(costs) == 110